#############################################################################
# Windows Wallpaper via MDM Personalization CSP
# Applies desktop and/or lock screen images via MDM Bridge Provider
# Runs safely under SYSTEM context (no FileWave API interaction)
#############################################################################

# Step 0 - Force 64-bit PowerShell if running under 32-bit on x64 systems
if ($env:PROCESSOR_ARCHITEW6432 -eq "AMD64") {
    if ($myInvocation.Line) {
        &"$env:WINDIR\sysnative\WindowsPowerShell\v1.0\powershell.exe" -NonInteractive -NoProfile $myInvocation.Line
    } else {
        &"$env:WINDIR\sysnative\WindowsPowerShell\v1.0\powershell.exe" -NonInteractive -NoProfile -File "$($myInvocation.InvocationName)" $args
    }
    exit $LASTEXITCODE
}

# --------------------------
# Configuration and Helpers
# --------------------------

function Write-Log([string]$Message, [string]$Level = 'INFO') {
    $ts = Get-Date -Format "s"
    Write-Host "[$ts] [$Level] $Message"
}

function Convert-ToFileUri {
    param([Parameter(Mandatory = $true)][string]$Path)
    $expanded = [Environment]::ExpandEnvironmentVariables($Path)

    # Add drive if starting with / or \
    if ($expanded -match '^[\\/]' -and -not ($expanded -match '^[A-Za-z]:')) {
        $expanded = "C:$expanded"
    }

    # If already a URI
    if ($expanded -match '^(?i)(file|http|https)://') { return $expanded }

    # Convert local path to file:///
    if ($expanded -match '^(?i)[A-Za-z]:[\\/]' -or $expanded -match '^[\\/]{2}') {
        $normalized = $expanded -replace '/', '\'
        if (-not (Test-Path -LiteralPath $normalized)) {
            throw "Local image not found: $normalized"
        }
        $full = (Resolve-Path -LiteralPath $normalized).Path
        $uri = New-Object System.Uri($full)
        return $uri.AbsoluteUri
    }

    throw "Unsupported image value: $Path (must be http(s):// or a valid local path)."
}

# --------------------------
# Input normalization
# --------------------------

$DesktopRaw = $Env:DESKTOP_IMAGE_URL
$LockRaw    = $Env:LOCKSCREEN_IMAGE_URL

if ([string]::IsNullOrWhiteSpace($DesktopRaw) -and [string]::IsNullOrWhiteSpace($LockRaw)) {
    Write-Log "Provide DESKTOP_IMAGE_URL and/or LOCKSCREEN_IMAGE_URL." 'ERROR'
    exit 2
}

try {
    if ($DesktopRaw) {
        $DesktopUrl = Convert-ToFileUri -Path $DesktopRaw
        Write-Log "Normalized DesktopImageUrl -> $DesktopUrl"
    }
    if ($LockRaw) {
        $LockUrl = Convert-ToFileUri -Path $LockRaw
        Write-Log "Normalized LockScreenImageUrl -> $LockUrl"
    }
} catch {
    Write-Log $_.Exception.Message 'ERROR'
    exit 3
}

# --------------------------
# Step 1 - Apply via MDM Bridge Provider
# --------------------------

$Namespace = 'root\cimv2\mdm\dmmap'
$Class     = 'MDM_Personalization'

try {
    $Instance = Get-CimInstance -Namespace $Namespace -ClassName $Class -ErrorAction Stop
} catch {
    Write-Log "Failed to access MDM Bridge Provider in '$Namespace'. Run as SYSTEM. $_" 'ERROR'
    exit 4
}

$Changed = $false
if ($DesktopUrl) {
    Write-Log "Setting DesktopImageUrl: $DesktopUrl"
    $Instance.DesktopImageUrl = $DesktopUrl
    $Changed = $true
}
if ($LockUrl) {
    Write-Log "Setting LockScreenImageUrl: $LockUrl"
    $Instance.LockScreenImageUrl = $LockUrl
    $Changed = $true
}

if ($Changed) {
    try {
        Set-CimInstance -CimInstance $Instance -ErrorAction Stop | Out-Null
        Write-Log "CSP values committed via WMI Bridge."
    } catch {
        Write-Log "Failed to write CSP values via WMI Bridge. $_" 'ERROR'
        exit 5
    }
} else {
    Write-Log "No values to change." 'WARN'
}

# --------------------------
# Step 2 - Poll CSP status for progress
# --------------------------

function Get-StatusText([int]$code) {
    switch ($code) {
        1 {'Success'}
        2 {'InProgress'}
        3 {'Failed'}
        4 {'UnknownFileType'}
        5 {'UnsupportedUrlScheme'}
        6 {'MaxRetryFailed'}
        default {'Unknown'}
    }
}

$Attempts = 12
for ($i = 1; $i -le $Attempts; $i++) {
    Start-Sleep -Seconds 10
    $Instance = Get-CimInstance -Namespace $Namespace -ClassName $Class
    $dStatus = $Instance.DesktopImageStatus
    $lStatus = $Instance.LockScreenImageStatus
    if ($DesktopUrl) { Write-Log ("DesktopImageStatus=$dStatus ({0})" -f (Get-StatusText $dStatus)) }
    if ($LockUrl)    { Write-Log ("LockScreenImageStatus=$lStatus ({0})" -f (Get-StatusText $lStatus)) }

    $doneDesktop = -not $DesktopUrl -or ($dStatus -in 1,3,4,5,6)
    $doneLock    = -not $LockUrl    -or ($lStatus -in 1,3,4,5,6)
    if ($doneDesktop -and $doneLock) { break }
}

Write-Log "Wallpaper configuration completed."
exit 0
